import sys
import re
import pandas as pd
import logging
from multiprocessing import Pool, cpu_count

# 配置日志格式
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# Excel文件名和主频
xlsx_filename = 'output.xlsx'
chip_frequency = 300_000_000  # 300 MHz

# 跳转指令列表
jump_instructions = ['JMP', 'JEQ', 'JLEZ', 'JNE', 'JNZ', 'JGEZ', 'JLTZ', 'JZ', 'JGE', 'JLT', 'J', 'JI', 'JGTZ', 'JNED', 'CALL']

# 定义汇编指令及其时钟周期
instruction_cycles = {
    'FMAX': 4,       # Maximum of floating-point numbers
    'FLOOR': 5,      # Floor value of floating-point number
    'TRUNC': 4,      # Truncate floating-point number
    'FMOD': 6,       # Floating-point modulo
    'FMAXF': 4,      # Maximum of single-precision floating-point numbers
    'FMINF': 4,      # Minimum of single-precision floating-point numbers
    'SQRTF': 5,      # Square root of single-precision floating-point number
    'FLOORF': 5,     # Floor value of single-precision floating-point number
    'FMODF': 6,      # Floating-point modulo of single-precision number
}

def is_jump_instruction(instruction):
    """
    判断指令是否为跳转指令。
    
    Args:
        instruction (str): 汇编指令字符串。
    
    Returns:
        bool: 如果是跳转指令返回True，否则返回False。
    """
    for jump_instr in jump_instructions:
        if instruction.upper().startswith(jump_instr):
            return True
    return False

def get_jump_target(instruction):
    """
    获取跳转指令的目标函数。
    
    Args:
        instruction (str): 汇编指令字符串。
    
    Returns:
        str: 目标函数名，如果没有找到则返回空字符串。
    """
    # 构建正则表达式
    jump_instr_pattern = '|'.join(jump_instructions)
    match = re.match(fr'({jump_instr_pattern})\s+(\w+)', instruction, re.IGNORECASE)
    if match:
        return match.group(2)
    return ''

def load_excel_data(filename):
    """
    使用 pandas 一次性读取 Excel 文件中的数据，并返回包含所有行的列表。
    
    Args:
        filename (str): Excel 文件名。
    
    Returns:
        list: 包含所有行的列表，每行是一个包含单元格值的元组。
    """
    df = pd.read_excel(filename)
    data = df.values.tolist()
    return data

def get_function_cycles(function_name, visited_functions, cycle_cache, data):
    """
    获取指定函数及其跳转函数的总时钟周期数。
    
    Args:
        function_name (str): 要查找的函数名。
        visited_functions (set): 已访问过的函数名集合，防止循环调用。
        cycle_cache (dict): 已计算过的函数时钟周期缓存。
        data (list): 从 Excel 文件中读取的数据。
    
    Returns:
        int: 函数全生命周期的总时钟周期数。
    """
    if function_name in visited_functions:
        return 0
    if function_name in cycle_cache:
        return cycle_cache[function_name]

    visited_functions.add(function_name)
    total_cycles = 0
    jump_functions = []

    # 查找函数的总时钟周期
    for row in data:
        if row[1] == function_name:
            if row[4]:
                total_cycles = row[4]
            if row[2] and is_jump_instruction(row[2]):
                jump_target = get_jump_target(row[2])
                if jump_target and jump_target != function_name:
                    jump_functions.append(jump_target)

    # 递归查找跳转函数的总时钟周期
    for jump_function in jump_functions:
        logging.info(f'Function {function_name} jumps to {jump_function}')
        total_cycles += get_function_cycles(jump_function, visited_functions, cycle_cache, data)

    cycle_cache[function_name] = total_cycles
    if total_cycles > 0:
        logging.info(f'Total cycles for function {function_name}: {total_cycles}')
    else:
        # 在字典中查找匹配项
        matched_cycle = instruction_cycles.get(function_name.upper(), None)
        if matched_cycle is not None:
            total_cycles = matched_cycle
            logging.info(f'Total cycles for function {function_name}: {total_cycles}')
        else:
            logging.warning(f'Total cycles for function {function_name}: {total_cycles}')
    return total_cycles

def main():
    if len(sys.argv) != 2:
        print("Usage: python xxx.py <function_name>")
        sys.exit(1)

    function_name = sys.argv[1]
    visited_functions = set()
    cycle_cache = {}

    # 使用 pandas 一次性读取 Excel 文件中的数据
    data = load_excel_data(xlsx_filename)

    # 使用多进程处理数据
    with Pool(cpu_count()) as pool:
        total_cycles = get_function_cycles(function_name, visited_functions, cycle_cache, data)

    # 计算运行时间，单位为毫秒
    execution_time_ms = (total_cycles / chip_frequency) * 1000

    print(f'Function {function_name} total lifecycle cycles: {total_cycles}')
    print(f'Function {function_name} execution time: {execution_time_ms:.6f} ms')

if __name__ == "__main__":
    main()
